# TODO: 导入必要的库和模块
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
import pickle
from tqdm import tqdm

# TODO: 加载数字数据集
# scikit-learn 自带了一个手写数字数据集，我们可以使用 load_digits() 函数加载它
digits = load_digits()
# 数据集中的图像数据，每个图像是一个 8x8 的像素矩阵，被展平成一个 64 维的向量
X = digits.data
# 数据集中的标签数据，即每个图像对应的数字
y = digits.target

# TODO: 将数据集划分为训练集和测试集
# 使用 train_test_split 函数将数据集随机划分为训练集和测试集
# test_size=0.25 表示将 25% 的数据用作测试集
# random_state=42 确保每次划分的结果都一样，便于复现
(X_train, X_test, y_train, y_test) = train_test_split(X, y, test_size=0.25, random_state=42)

# TODO: 初始化变量以存储最佳准确率，相应的k值和最佳knn模型
best_accuracy = 0.0
best_k = -1
best_knn = None

# TODO: 初始化一个列表以存储每个k值的准确率
# 这个列表可以用来后续分析不同k值对准确率的影响
accuracies = []

# TODO: 尝试从1到40的k值，对于每个k值，训练knn模型，保存最佳准确率，k值和knn模型
# 我们将测试从 1 到 40 的 k 值，找出哪个 k 值能带来最高的模型准确率
# 使用 tqdm 库可以显示一个进度条，方便我们查看进度
print("正在寻找最佳 k 值...")
for k in tqdm(range(1, 41)):
    # 初始化KNeighborsClassifier模型，设置邻居数量为当前的k值
    knn = KNeighborsClassifier(n_neighbors=k)
    
    # 使用训练数据来训练模型
    knn.fit(X_train, y_train)
    
    # 使用测试数据来评估模型的准确率
    accuracy = knn.score(X_test, y_test)
    
    # 将当前k值和对应的准确率存入列表
    accuracies.append(accuracy)
    
    # 如果当前模型的准确率比之前记录的最高准确率还要高
    if accuracy > best_accuracy:
        # 更新最高准确率
        best_accuracy = accuracy
        # 更新最佳k值
        best_k = k
        # 更新最佳模型
        best_knn = knn

# TODO: 将最佳KNN模型保存到二进制文件
# 使用 pickle 库将训练好的最佳模型对象序列化并保存到文件中
# 'wb' 表示以二进制写模式打开文件
print(f"\n训练完成。正在将最佳模型保存到 'best_knn_model.pkl'...")
with open('best_knn_model.pkl', 'wb') as f:
    pickle.dump(best_knn, f)
print("模型保存成功。")

# TODO: 打印最佳准确率和相应的k值
print(f"最佳 k 值: {best_k}")
print(f"对应的测试集准确率: {best_accuracy:.4f}")